'''
霍尼项目
训练指定点位代码
使用方法:
1 标注文件放到add_imgs目录下
2 运行下述命令，训练纸盒模型, 模型保存至'D:\data\231207huoni\trainseg_{name}\models'
cd Yolov5_insSeg
python trainseg_auto.py --name zhihe

'''

import argparse
import datetime
import glob
import hashlib
import os
import sys

import yaml

from yolov5.myutil_seg import labelme2yoloseg
# from pathlib import Path

# from yolov5.utils.general import increment_path # exp1, exp2, ...
sys.path.append('yolov5')
from yolov5.utils.dataloaders import get_hash # 避免重复转换数据集



def getn(save_path): # 获取expn
    li = os.listdir(save_path) # models目录得是exp,exp2,...
    inds = [-1]
    for i in li:
        if not i[:3] == 'exp':
            continue
        p = i[3:]
        inds.append(0) if p == '' else inds.append(int(p))

    n = max(inds)
    # if n == 0:
    #     n = ''

    return n


def trainseg_auto(name ='zhihe', root_path = r'D:\data\231207huoni',random_scale=1.0):


    # 1 转换格式 labelme转为yolo数据集
    l2y = labelme2yoloseg()
    l2y.jpg = fr'{root_path}/trainseg_{name}/add_imgs/*/*.jpg'
    l2y.json = fr'{root_path}/trainseg_{name}/add_imgs/*/*.json'
    l2y.txt = fr'{root_path}/trainseg_{name}/add_imgs/*/*.txt'
    l2y.cls = fr'{root_path}/trainseg_{name}/add_imgs/classes.txt'
    # l2y.cls = os.path.join(root_path, f'trainseg_{name}', 'add_imgs','classes.txt')  #
    l2y.md5 = fr'{root_path}/trainseg_{name}/add_imgs/md5.txt' # jpg和json文件
    l2y.target_size = 640  # (2048, 3072, 3)

    l2y.yolov5 = fr'{root_path}/trainseg_{name}/imglabs_yolov5/' # yolov5数据集地址
    l2y.buff = datetime.datetime.now().strftime('d%Y%m%d_') # 前缀

    # 读取label
    l2y.ref = {}
    # l2y.ref = {'zhihe': 0, 'logo': 1}  #
    with open(l2y.cls) as f:
        lbs = f.read().strip().splitlines()
    for ind, lb in enumerate(lbs):
        l2y.ref[lb.strip()] = ind

    # 转换
    isUpdate = False  # 是否转换数据集
    jpg_paths = glob.glob(fr'{root_path}/trainseg_{name}/add_imgs/*[!mini]/*.jpg') # 不包含mini数据
    json_paths = glob.glob(fr'{root_path}/trainseg_{name}/add_imgs/*[!mini]/*.json')
    cur_md5_value = get_hash(jpg_paths + json_paths)
    if not os.path.exists(l2y.md5):
        with open(l2y.md5, 'w') as fp:
            fp.write('')
        # os.mknod(l2y.md5) # wind
    with open(l2y.md5, 'r') as fp:
        pre_md5_Value = fp.read()
    if not pre_md5_Value == cur_md5_value:
        with open(l2y.md5, 'w') as fp:
            fp.write(cur_md5_value)
        isUpdate = True
    print(f'（不含_mini文件夹）原始图片和json文件是否变化 {isUpdate}')
    if isUpdate:
        # l2y.run(isAug = False) # 转换
        l2y.run(isAug=True, random_scale=random_scale)  # 转换

    # 2 写入huoni_seg_{name}.yaml
    data_yaml_path = f'yolov5/data/huoni_seg_{name}.yaml'
    # with open(data_yaml_path, errors='ignore') as f:
    #     data_yaml = yaml.safe_load(f)
    data_yaml = {}
    data_yaml['nc'] = len(lbs)
    data_yaml['names'] = lbs
    data_yaml['train'] = [f'{l2y.yolov5}images/train']
    data_yaml['val'] = [f'{l2y.yolov5}images/val']
    with open(data_yaml_path, 'w', encoding='utf-8') as f:
        yaml.dump(data=data_yaml, stream=f, allow_unicode=True) # 保存

    # 3 训练
    print('\n\n\n')
    save_path = fr'{root_path}/trainseg_{name}/models' # 模型保存路径
    hpy = 'yolov5/data/hyps/hyp.scratch-low-huoni.yaml' # 超参
    imgsz = l2y.target_size # 416

    n = getn(save_path)
    weights = 'yolov5/yolov5s-seg.pt'
    if n != -1:
        if n == 0:
            n = ''
        weights = fr'{save_path}/exp{n}/weights/best.pt'
    train_cmd_str = (f'python yolov5/segment/train.py --epoch {300} --weights {weights} --data {data_yaml_path} --imgsz {imgsz} '
                     f'--batch-size {4} --project {save_path} --hyp {hpy} --no-overlap '
                     f'--mask-ratio {1} --patience {60} --device {1}' )
    print(train_cmd_str)
    os.system(train_cmd_str)


    # 4 测试
    # for n in range(1, 9999):
    #     if n == 1:
    #         p = f'{save_path}\exp'
    #     else:
    #         p = f'{save_path}\exp{n}'  # increment path
    #     if not os.path.exists(p):  #
    #         model_path = fr'{save_path}\exp{n}\weights\best.pt'
    #         break
    li = os.listdir(save_path) # models目录得是exp,exp2,...
    inds = []
    for i in li:
        if not i[:3] == 'exp':
            continue
        p = i[3:]
        inds.append(0) if p == '' else inds.append(int(p))
    n = max(inds)
    if n == 0:
        n = ''
    model_path = fr'{save_path}/exp{n}/weights/best.pt' # 测试模型地址
    out_test_path = fr'{root_path}/trainseg_{name}/predict' # 测试结果图片
    pre_cmd_str = f'python yolov5/segment/predict.py --weights {model_path} --source {l2y.jpg}  --project {out_test_path}'
    print(pre_cmd_str)
    os.system(pre_cmd_str)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument('--name', type=str, default='strap', help='点位名称，或物料名称')

    parser.add_argument('--root_path', type=str, default=r'/home/ps/zhangxiancai/data/231207huoni', help='数据集总目录，下一级为trainseg_zhihe等')
    parser.add_argument('--random_scale', type=float, default=2.0,
                        help='自定义数据增强：随机裁剪框大小，原框倍数')
    par = parser.parse_args()
    # name = par.name
    # root_path = par.root_path

    names = par.name.split(' ')
    print(names)
    for name in names:
        trainseg_auto(name, par.root_path, par.random_scale)
        # try:
        #     trainseg_auto(name, par.root_path)
        # except Exception as e:
        #     print(e)